from starlette.applications import Starlette
from starlette.responses import HTMLResponse
from starlette.staticfiles import StaticFiles
from starlette.middleware.cors import CORSMiddleware
import uvicorn, aiohttp, asyncio
from io import BytesIO
from fastai.vision import *
from gradcam import *
import base64

model_file_url= 'https://drive.google.com/uc?export=download&id=1_BkMp1DwT_8QwFogLRidSGFDTY9mAaoD'
model_file_name = 'model'
classes = ['beauty','castle','hercules','howl','kiki','mermaid','mononoke','mulan','pocahontas','tarzan']
mv_names= ['Beauty and the Beast',
            'Laputa Castle in the Sky',
            'Hercules',
            'Howl''s Moving Castle',
            'Kiki''s Delivery Service',
            'The Little Mermaid',
            'Princess Mononoke',
            'Mulan',
            'Pocahontas',
            'Tarzan']
mv_dict = {i:v for i,v in zip(classes,mv_names)}

path = Path(__file__).parent

app = Starlette()
app.add_middleware(CORSMiddleware, allow_origins=['*'], allow_headers=['X-Requested-With', 'Content-Type'])
app.mount('/static', StaticFiles(directory='app/static'))
app.mount('/tmp', StaticFiles(directory='/tmp'))

async def download_file(url, dest):
    if dest.exists(): return
    async with aiohttp.ClientSession() as session:
        async with session.get(url) as response:
            data = await response.read()
            with open(dest, 'wb') as f: f.write(data)

async def setup_learner():
    await download_file(model_file_url, path/'models'/f'{model_file_name}.pth')
    data_bunch = ImageDataBunch.single_from_classes(path, classes,
        tfms=get_transforms(), size=350).normalize(imagenet_stats)
    learn = create_cnn(data_bunch, models.resnet50, pretrained=False)
    learn.load(model_file_name)
    return learn

loop = asyncio.get_event_loop()
tasks = [asyncio.ensure_future(setup_learner())]
learn = loop.run_until_complete(asyncio.gather(*tasks))[0]
loop.close()

@app.route("/upload", methods=["POST"])
async def upload(request):
    data = await request.form()
    img_bytes = await (data["file"].read())
    # bytes = base64.b64decode(img_bytes)
    return predict_from_bytes(img_bytes)

def predict_from_bytes(bytes):
    img = open_image(BytesIO(bytes))
    _,_,losses = learn.predict(img)
    predictions = sorted(zip(classes, map(float, losses)), key=lambda p: p[1], reverse=True)

    rs = '<p>Top 3 predictions:</p>\n'
    for clas,pr in predictions[:3]:
        rs+=f'<p> -{mv_dict[clas]}: {(pr*100):.2f}% </p>\n'
    if predictions[0][1] <= 0.70:
        rs+='<p>(Note: Model is not confident with this prediction)</p>\n'

    rs+=f'<p>Which part of the image the model considered for <b>{mv_dict[predictions[0][0]]}</b> prediction: </p>\n'
    
    gcam = GradCam.from_one_img(learn,img)
    gcam.plot();

    result_html1 = path/'static'/'result1.html'
    result_html2 = path/'static'/'result2.html'
    
    result_html = str(result_html1.open().read() +rs + result_html2.open().read())
    return HTMLResponse(result_html)

@app.route("/")
def form(request):
    index_html = path/'static'/'index.html'
    return HTMLResponse(index_html.open().read())

if __name__ == "__main__":
    if "serve" in sys.argv: uvicorn.run(app, host="0.0.0.0", port=8080)
